
!pip install tensorflow==2.10 wandb

Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/
Collecting tensorflow==2.10
  Downloading tensorflow-2.10.0-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (578.1 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 578.1/578.1 MB 1.9 MB/s eta 0:00:00
Collecting wandb
  Downloading wandb-0.13.9-py2.py3-none-any.whl (2.0 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 2.0/2.0 MB 33.0 MB/s eta 0:00:00
Collecting tensorflow-estimator<2.11,>=2.10.0
  Downloading tensorflow_estimator-2.10.0-py2.py3-none-any.whl (438 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 438.7/438.7 KB 21.6 MB/s eta 0:00:00
Requirement already satisfied: grpcio<2.0,>=1.24.3 in /usr/local/lib/python3.8/dist-packages (from tensorflow==2.10) (1.51.1)
Requirement already satisfied: termcolor>=1.1.0 in /usr/local/lib/python3.8/dist-packages (from tensorflow==2.10) (2.2.0)
Requirement already satisfied: numpy>=1.20 in /usr/local/lib/python3.8/dist-packages (from tensorflow==2.10) (1.21.6)
Requirement already satisfied: protobuf<3.20,>=3.9.2 in /usr/local/lib/python3.8/dist-packages (from tensorflow==2.10) (3.19.6)
Requirement already satisfied: absl-py>=1.0.0 in /usr/local/lib/python3.8/dist-packages (from tensorflow==2.10) (1.3.0)
Requirement already satisfied: keras-preprocessing>=1.1.1 in /usr/local/lib/python3.8/dist-packages (from tensorflow==2.10) (1.1.2)
Requirement already satisfied: packaging in /usr/local/lib/python3.8/dist-packages (from tensorflow==2.10) (21.3)
Requirement already satisfied: astunparse>=1.6.0 in /usr/local/lib/python3.8/dist-packages (from tensorflow==2.10) (1.6.3)
Requirement already satisfied: tensorflow-io-gcs-filesystem>=0.23.1 in /usr/local/lib/python3.8/dist-packages (from tensorflow==2.10) (0.29.0)
Requirement already satisfied: gast<=0.4.0,>=0.2.1 in /usr/local/lib/python3.8/dist-packages (from tensorflow==2.10) (0.4.0)
Requirement already satisfied: h5py>=2.9.0 in /usr/local/lib/python3.8/dist-packages (from tensorflow==2.10) (3.1.0)
Requirement already satisfied: setuptools in /usr/local/lib/python3.8/dist-packages (from tensorflow==2.10) (57.4.0)
Collecting flatbuffers>=2.0
  Downloading flatbuffers-23.1.21-py2.py3-none-any.whl (26 kB)
Requirement already satisfied: typing-extensions>=3.6.6 in /usr/local/lib/python3.8/dist-packages (from tensorflow==2.10) (4.4.0)
Requirement already satisfied: six>=1.12.0 in /usr/local/lib/python3.8/dist-packages (from tensorflow==2.10) (1.15.0)
Requirement already satisfied: opt-einsum>=2.3.2 in /usr/local/lib/python3.8/dist-packages (from tensorflow==2.10) (3.3.0)
Collecting keras<2.11,>=2.10.0
  Downloading keras-2.10.0-py2.py3-none-any.whl (1.7 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 1.7/1.7 MB 13.8 MB/s eta 0:00:00
Requirement already satisfied: wrapt>=1.11.0 in /usr/local/lib/python3.8/dist-packages (from tensorflow==2.10) (1.14.1)
Requirement already satisfied: libclang>=13.0.0 in /usr/local/lib/python3.8/dist-packages (from tensorflow==2.10) (15.0.6.1)
Requirement already satisfied: google-pasta>=0.1.1 in /usr/local/lib/python3.8/dist-packages (from tensorflow==2.10) (0.2.0)
Collecting tensorboard<2.11,>=2.10
  Downloading tensorboard-2.10.1-py3-none-any.whl (5.9 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 5.9/5.9 MB 59.5 MB/s eta 0:00:00
Requirement already satisfied: appdirs>=1.4.3 in /usr/local/lib/python3.8/dist-packages (from wandb) (1.4.4)
Requirement already satisfied: PyYAML in /usr/local/lib/python3.8/dist-packages (from wandb) (6.0)
Requirement already satisfied: psutil>=5.0.0 in /usr/local/lib/python3.8/dist-packages (from wandb) (5.4.8)
Collecting docker-pycreds>=0.4.0
  Downloading docker_pycreds-0.4.0-py2.py3-none-any.whl (9.0 kB)
Collecting pathtools
  Downloading pathtools-0.1.2.tar.gz (11 kB)
  Preparing metadata (setup.py) ... done
Requirement already satisfied: Click!=8.0.0,>=7.0 in /usr/local/lib/python3.8/dist-packages (from wandb) (7.1.2)
Requirement already satisfied: requests<3,>=2.0.0 in /usr/local/lib/python3.8/dist-packages (from wandb) (2.25.1)
Collecting setproctitle
  Downloading setproctitle-1.3.2-cp38-cp38-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl (31 kB)
Collecting sentry-sdk>=1.0.0
  Downloading sentry_sdk-1.14.0-py2.py3-none-any.whl (178 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 178.9/178.9 KB 22.5 MB/s eta 0:00:00
Collecting GitPython>=1.0.0
  Downloading GitPython-3.1.30-py3-none-any.whl (184 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 184.0/184.0 KB 20.2 MB/s eta 0:00:00
Requirement already satisfied: wheel<1.0,>=0.23.0 in /usr/local/lib/python3.8/dist-packages (from astunparse>=1.6.0->tensorflow==2.10) (0.38.4)
Collecting gitdb<5,>=4.0.1
  Downloading gitdb-4.0.10-py3-none-any.whl (62 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 62.7/62.7 KB 7.8 MB/s eta 0:00:00
Requirement already satisfied: chardet<5,>=3.0.2 in /usr/local/lib/python3.8/dist-packages (from requests<3,>=2.0.0->wandb) (4.0.0)
Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.8/dist-packages (from requests<3,>=2.0.0->wandb) (2.10)
Requirement already satisfied: urllib3<1.27,>=1.21.1 in /usr/local/lib/python3.8/dist-packages (from requests<3,>=2.0.0->wandb) (1.24.3)
Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.8/dist-packages (from requests<3,>=2.0.0->wandb) (2022.12.7)
Collecting urllib3<1.27,>=1.21.1
  Downloading urllib3-1.26.14-py2.py3-none-any.whl (140 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 140.6/140.6 KB 8.9 MB/s eta 0:00:00
Requirement already satisfied: markdown>=2.6.8 in /usr/local/lib/python3.8/dist-packages (from tensorboard<2.11,>=2.10->tensorflow==2.10) (3.4.1)
Requirement already satisfied: tensorboard-data-server<0.7.0,>=0.6.0 in /usr/local/lib/python3.8/dist-packages (from tensorboard<2.11,>=2.10->tensorflow==2.10) (0.6.1)
Requirement already satisfied: werkzeug>=1.0.1 in /usr/local/lib/python3.8/dist-packages (from tensorboard<2.11,>=2.10->tensorflow==2.10) (1.0.1)
Requirement already satisfied: google-auth-oauthlib<0.5,>=0.4.1 in /usr/local/lib/python3.8/dist-packages (from tensorboard<2.11,>=2.10->tensorflow==2.10) (0.4.6)
Requirement already satisfied: google-auth<3,>=1.6.3 in /usr/local/lib/python3.8/dist-packages (from tensorboard<2.11,>=2.10->tensorflow==2.10) (2.16.0)
Requirement already satisfied: tensorboard-plugin-wit>=1.6.0 in /usr/local/lib/python3.8/dist-packages (from tensorboard<2.11,>=2.10->tensorflow==2.10) (1.8.1)
Requirement already satisfied: pyparsing!=3.0.5,>=2.0.2 in /usr/local/lib/python3.8/dist-packages (from packaging->tensorflow==2.10) (3.0.9)
Collecting smmap<6,>=3.0.1
  Downloading smmap-5.0.0-py3-none-any.whl (24 kB)
Requirement already satisfied: pyasn1-modules>=0.2.1 in /usr/local/lib/python3.8/dist-packages (from google-auth<3,>=1.6.3->tensorboard<2.11,>=2.10->tensorflow==2.10) (0.2.8)
Requirement already satisfied: rsa<5,>=3.1.4 in /usr/local/lib/python3.8/dist-packages (from google-auth<3,>=1.6.3->tensorboard<2.11,>=2.10->tensorflow==2.10) (4.9)
Requirement already satisfied: cachetools<6.0,>=2.0.0 in /usr/local/lib/python3.8/dist-packages (from google-auth<3,>=1.6.3->tensorboard<2.11,>=2.10->tensorflow==2.10) (5.2.1)
Requirement already satisfied: requests-oauthlib>=0.7.0 in /usr/local/lib/python3.8/dist-packages (from google-auth-oauthlib<0.5,>=0.4.1->tensorboard<2.11,>=2.10->tensorflow==2.10) (1.3.1)
Requirement already satisfied: importlib-metadata>=4.4 in /usr/local/lib/python3.8/dist-packages (from markdown>=2.6.8->tensorboard<2.11,>=2.10->tensorflow==2.10) (6.0.0)
Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.8/dist-packages (from importlib-metadata>=4.4->markdown>=2.6.8->tensorboard<2.11,>=2.10->tensorflow==2.10) (3.11.0)
Requirement already satisfied: pyasn1<0.5.0,>=0.4.6 in /usr/local/lib/python3.8/dist-packages (from pyasn1-modules>=0.2.1->google-auth<3,>=1.6.3->tensorboard<2.11,>=2.10->tensorflow==2.10) (0.4.8)
Requirement already satisfied: oauthlib>=3.0.0 in /usr/local/lib/python3.8/dist-packages (from requests-oauthlib>=0.7.0->google-auth-oauthlib<0.5,>=0.4.1->tensorboard<2.11,>=2.10->tensorflow==2.10) (3.2.2)
Building wheels for collected packages: pathtools
  Building wheel for pathtools (setup.py) ... done
  Created wheel for pathtools: filename=pathtools-0.1.2-py3-none-any.whl size=8806 sha256=4c49843404693307a7bf620034caf681be91c98ed9cd1154db5557727bf8e879
  Stored in directory: /root/.cache/pip/wheels/4c/8e/7e/72fbc243e1aeecae64a96875432e70d4e92f3d2d18123be004
Successfully built pathtools
Installing collected packages: pathtools, keras, flatbuffers, urllib3, tensorflow-estimator, smmap, setproctitle, docker-pycreds, sentry-sdk, gitdb, GitPython, wandb, tensorboard, tensorflow
  Attempting uninstall: keras
    Found existing installation: keras 2.9.0
    Uninstalling keras-2.9.0:
      Successfully uninstalled keras-2.9.0
  Attempting uninstall: flatbuffers
    Found existing installation: flatbuffers 1.12
    Uninstalling flatbuffers-1.12:
      Successfully uninstalled flatbuffers-1.12
  Attempting uninstall: urllib3
    Found existing installation: urllib3 1.24.3
    Uninstalling urllib3-1.24.3:
      Successfully uninstalled urllib3-1.24.3
  Attempting uninstall: tensorflow-estimator
    Found existing installation: tensorflow-estimator 2.9.0
    Uninstalling tensorflow-estimator-2.9.0:
      Successfully uninstalled tensorflow-estimator-2.9.0
  Attempting uninstall: tensorboard
    Found existing installation: tensorboard 2.9.1
    Uninstalling tensorboard-2.9.1:
      Successfully uninstalled tensorboard-2.9.1
  Attempting uninstall: tensorflow
    Found existing installation: tensorflow 2.9.2
    Uninstalling tensorflow-2.9.2:
      Successfully uninstalled tensorflow-2.9.2
Successfully installed GitPython-3.1.30 docker-pycreds-0.4.0 flatbuffers-23.1.21 gitdb-4.0.10 keras-2.10.0 pathtools-0.1.2 sentry-sdk-1.14.0 setproctitle-1.3.2 smmap-5.0.0 tensorboard-2.10.1 tensorflow-2.10.0 tensorflow-estimator-2.10.0 urllib3-1.26.14 wandb-0.13.9


from tensorflow.keras.layers import LeakyReLU, Conv2D, Dropout , Concatenate ,UpSampling2D, SpatialDropout2D , GlobalAveragePooling2D  , Normalization, BatchNormalization , Activation   , AveragePooling2D  , MaxPooling2D, ZeroPadding2D , Dense , Add , ZeroPadding2D  , Input ,Flatten , Reshape
from tensorflow.keras.models import Model , Sequential
import tensorflow as tf 
import tensorflow.keras.backend as K 
from tensorflow.keras import mixed_precision
from tensorflow.keras.regularizers import L2
from tensorflow.keras.callbacks import ReduceLROnPlateau
import gc
from sklearn.model_selection import train_test_split
# mixed_precision.set_global_policy('mixed_float16')
IMG_SIZE = 28
import matplotlib.pyplot as plt
tf.random.set_seed(42)


from tensorflow.python.util.nest import pack_sequence_as
import tensorflow.keras.models as KM
import tensorflow as tf

import wandb
import warnings

from tensorflow.keras.applications.inception_v3 import InceptionV3
from tensorflow.keras import backend as K
import numpy as np

class KID(keras.metrics.Metric):
	#https://keras.io/examples/generative/gan_ada/ 
    def __init__(self, name="kid", **kwargs):
        super().__init__(name=name, **kwargs)

        # KID is estimated per batch and is averaged across batches
        self.kid_tracker = keras.metrics.Mean()

        # a pretrained InceptionV3 is used without its classification layer
        # transform the pixel values to the 0-255 range, then use the same
        # preprocessing as during pretraining
        self.encoder = keras.Sequential(
            [
                layers.InputLayer(input_shape=(32, 32, 3)),
                # layers.Rescaling(255.0),
                layers.Resizing(height=75, width=75),
                layers.Lambda(keras.applications.inception_v3.preprocess_input),
                keras.applications.InceptionV3(
                    include_top=False,
                    input_shape=(75, 75, 3),
                    weights="imagenet",
                ),
                layers.GlobalAveragePooling2D(),
            ],
            name="inception_encoder",
        )

    def polynomial_kernel(self, features_1, features_2):
        feature_dimensions = tf.cast(tf.shape(features_1)[1], dtype=tf.float32)
        return (features_1 @ tf.transpose(features_2) / feature_dimensions + 1.0) ** 3.0

    def update_state(self, real_images, generated_images, sample_weight=None):
        real_features = self.encoder(real_images, training=False)
        generated_features = self.encoder(generated_images, training=False)

        # compute polynomial kernels using the two sets of features
        kernel_real = self.polynomial_kernel(real_features, real_features)
        kernel_generated = self.polynomial_kernel(
            generated_features, generated_features
        )
        kernel_cross = self.polynomial_kernel(real_features, generated_features)

        # estimate the squared maximum mean discrepancy using the average kernel values
        batch_size = tf.shape(real_features)[0]
        batch_size_f = tf.cast(batch_size, dtype=tf.float32)
        mean_kernel_real = tf.reduce_sum(kernel_real * (1.0 - tf.eye(batch_size))) / (
            batch_size_f * (batch_size_f - 1.0)
        )
        mean_kernel_generated = tf.reduce_sum(
            kernel_generated * (1.0 - tf.eye(batch_size))
        ) / (batch_size_f * (batch_size_f - 1.0))
        mean_kernel_cross = tf.reduce_mean(kernel_cross)
        kid = mean_kernel_real + mean_kernel_generated - 2.0 * mean_kernel_cross

        # update the average KID estimate
        self.kid_tracker.update_state(kid)

    def result(self):
        return self.kid_tracker.result()

    def reset_state(self):
        self.kid_tracker.reset_state()


class ImageLogger(tf.keras.callbacks.Callback):
	def __init__(self, num_row=6 , num_col = 10, latent_dim=128):
		self.num_rows = num_row
		self.latent_dim = latent_dim
		self.num_cols = num_col


	def on_epoch_end(self, epoch, logs=None):



		random_latent_vectors = tf.random.normal(shape=(self.num_rows * self.num_cols , self.latent_dim))
		generated_images = self.model.generator(random_latent_vectors).numpy() 
		generated_images = (generated_images +1 ) /2 
		plt.figure(figsize=(self.num_cols * 1.0, self.num_rows * 1.0))
		for row in range(self.num_rows):
			for col in range(self.num_cols):
				index = row * self.num_cols + col
				plt.subplot(self.num_rows, self.num_cols, index + 1)
				plt.imshow(generated_images[index] )
				plt.axis("off")
		plt.tight_layout()
		plt.show()
		plt.close()
import os 
import shutil
class WandbModelLogger(tf.keras.callbacks.Callback):
	def __init__(self, num_row=6 , num_col = 10, latent_dim=128):
		pack_sequence_as
		# self.generator_artifact = wandb.Artifact( name = wandb.run.name + ' Generator', type = 'model')
		# self.discriminator_artifact = wandb.Artifact( name = wandb.run.name + ' discriminator', type = 'model')
	



	def on_epoch_end(self, epoch, logs=None):
		if self.model.kid.name in logs:

			disstr = f"Discriminator-Epoch{epoch},KID={logs[self.model.kid.name]}"
			genstr = f"Generator-Epoch{epoch},KID={logs[self.model.kid.name]}"
		else:

			disstr = f"Discriminator-Epoch{epoch}"
			genstr = f"Generator-Epoch{epoch}"
		
		self.model.discriminator.save(disstr,save_format="tf")
		self.model.generator.save(genstr	, save_format="tf")
		# self.model.save(self.filepath[:-3], overwrite=True, save_format="tf")

        # Log the model as artifact.
		dis_name = wandb.util.make_artifact_name_safe(f"Discriminator-{wandb.run.name}")
		gen_name = wandb.util.make_artifact_name_safe(f"Generator-{wandb.run.name}")
		dis_artifact = wandb.Artifact(dis_name, type="model")
		gen_artifact = wandb.Artifact(gen_name, type="model")
  


		dis_artifact.add_dir(disstr)
		gen_artifact.add_dir(genstr)
		wandb.run.log_artifact(dis_artifact, aliases=["latest", f"Discriminator_epoch_{epoch}"])
		wandb.run.log_artifact(gen_artifact, aliases=["latest", f"Generator_epoch_{epoch}"])

		# Remove the SavedModel from wandb dir as we don't want to log it to save memory.
		shutil.rmtree(disstr)
		shutil.rmtree(genstr)
	# # import shutil
	# 	shutil.make_archive(disstr , 'zip', disstr)
	# 	shutil.make_archive(genstr , 'zip', genstr)
		
	
		
	# 	wandb.save( disstr + '.zip' , policy = 'now')
	# 	wandb.save(genstr+ '.zip' , policy = 'now')
	
	# 	shutil.rmtree(disstr )
	# 	shutil.rmtree(genstr)
	# 	os.remove( disstr+ '.zip')
	# 	os.remove(genstr+ '.zip')





#https://medium.com/dive-into-ml-ai/customization-of-model-fit-to-write-a-combined-discriminator-generator-gan-trainer-in-keras-524bce10cf66
class GAN(KM.Model):
	def __init__(self, discriminator, generator , latdim ,*args):
		super().__init__()
		self.discriminator = discriminator
		self.generator= generator
		self.latent_dim = latdim
		# self.val_data = val_data
		self.fid_history = dict(epoch = list() , fidscore = list() ) 
		self.test_step_counter = tf.Variable( 0, tf.int32)
  
	@property
	def metrics(self):
		return [self.kid]



	def compile(self, optimizers, loss):
		super().compile()

		# Optimizer for discriminator
		self.d_optimizer = optimizers["d_optimizer"]
		# Optimizer for generator
		self.g_optimizer = optimizers["g_optimizer"]

		# Loss for the main objective
		self.g_loss = loss["g_loss"]
		# Discriminator loss
		self.d_loss = loss["d_loss"]
		self.kid = KID()
		# self.fid = FID() 

		# self.metrics = metrics 

	def fit(self, *args, **kwargs):
		self.img = args[0] if len(args) >= 1 else kwargs['x']
		self.fid_history['epoch'] = list() 
		self.fid_history['fidscore'] = list() 
		self.test_step_per_epoch= tf.constant(kwargs['test_step_per_epoch'], tf.int32)
		del kwargs['test_step_per_epoch']
		return super().fit( *args, **kwargs)
  

	def test_step(self, real_images):
		if isinstance(real_images, tuple):
			real_images = real_images[0]
		
		batch_size = tf.shape(real_images)[0]
		latent_samples = tf.random.normal(shape=(batch_size, self.latent_dim))
		generated_images = self.generator(latent_samples, training=False)
     

		self.kid.update_state(real_images, (generated_images+1)/2)


		# only KID is measured during the evaluation phase for computational efficiency
		return {self.kid.name: self.kid.result() } 


	def train_step(self, real_images):
		if isinstance(real_images, tuple):
			real_images = real_images[0]
		# Sample random points in the latent space
		batch_size = tf.shape(real_images)[0]
		random_latent_vectors = tf.random.normal(shape=(batch_size, self.latent_dim))

		# Decode them to fake images
		generated_images = self.generator(random_latent_vectors , training = True)


		# Add random noise to the labels - important trick!
		# labels += 0.05 * tf.random.uniform(tf.shape(labels))

		# Train the discriminator
		def train_discriminator(label, img):
			with tf.GradientTape() as tape:
				predictions = self.discriminator(img, training = True)
				d_loss = self.d_loss(label, predictions)
				# d_loss_scaled = self.d_optimizer.get_scaled_loss(d_loss)
			grads = tape.gradient(d_loss, self.discriminator.trainable_weights)




			# grads = self.d_optimizer.get_unscaled_gradients( scaled_grads)
			self.d_optimizer.apply_gradients(
				zip(grads, self.discriminator.trainable_weights)
			)
            

			del predictions  , grads , tape , label
			return d_loss 
		d_loss_fake = train_discriminator(tf.ones((batch_size, 1)) +0.05 * tf.random.uniform((batch_size, 1)), generated_images)
		d_loss_real=train_discriminator(tf.zeros((batch_size, 1)) +0.05 * tf.random.uniform((batch_size, 1)), real_images)
		d_loss = 0.5*(d_loss_fake+d_loss_real)


		# Sample random points in the latent space


		random_latent_vectors = tf.random.normal(shape=(batch_size, self.latent_dim))

		# Assemble labels that say "all real images"
		misleading_labels = tf.zeros((batch_size, 1))

		# Train the generator
		# Do not update the discriminator weights

		with tf.GradientTape() as tape:
			predictions = self.discriminator(self.generator(random_latent_vectors, training = True), training = True)
			g_loss = self.g_loss(misleading_labels, predictions)
		grads  = tape.gradient(g_loss, self.generator.trainable_weights)
		self.g_optimizer.apply_gradients(zip(grads, self.generator.trainable_weights))
		del predictions, grads  , tape , misleading_labels 
		return {"d_loss": d_loss, "g_loss": g_loss}


def vggblock(x , n_conv, n_filters):
    for _ in range(n_conv-1 ):
        x = Conv2D(filters=n_filters,kernel_size=(3,3),padding="same")(x)
        x= BatchNormalization(momentum=0.1)(x)
        x = LeakyReLU(0.2)(x)
        

    x =  Conv2D(filters=n_filters,kernel_size=(3,3), strides = 2 ,padding="same" )(x)
    x= BatchNormalization(momentum=0.1)(x)
    x= LeakyReLU(0.2)(x)
    
    return x
def createcustomvgg_discriminator():#real vs fake classification
    inputtensor = Input(shape=(32,32,3))
    # x = vggblock(inputtensor, 2, 16)
    # x= MaxPooling2D((2,2) , strides =2 )(x)
  
    x = vggblock(inputtensor , 1 , 32)
    # x= MaxPooling2D((2,2) , strides =2 )(x)
    x = SpatialDropout2D(0.2)(x)
    x = vggblock(x , 1 , 64)
    # x =  MaxPooling2D((2,2) , strides =2 )(x)
    x = SpatialDropout2D(0.3)(x)
    x = vggblock(x , 1 , 128)
    # x = MaxPooling2D((2,2) , strides =2 )(x)
    x = SpatialDropout2D(0.3)(x)
    x = vggblock(x, 1, 256)
    x = GlobalAveragePooling2D()(x)
    # x = Dropout(0.5)(x)
    x= Dense(1 , 'sigmoid')(x)
    return Model( inputs = [inputtensor] , outputs = [x])



def gen_block(x, filters , n_conv, last  = False, kernel_size=(3, 3), strides=(1, 1), up_size=(2, 2)):
   
        # x = BatchNormalization()(x)

    
    x = UpSampling2D(up_size, interpolation='bilinear')(x)
    x = Conv2D(filters, kernel_size, strides=strides, padding='same', use_bias=False )(x)
    x= BatchNormalization(momentum=0.1)(x)
    x =  tf.keras.layers.ReLU()(x)
    for _ in range(n_conv-1):
        x = Conv2D(filters, kernel_size, strides=strides, padding='same', use_bias=False )(x)
        x= BatchNormalization(momentum=0.1)(x) 
        x =  tf.keras.layers.ReLU()(x)
       

    
    # x = BatchNormalization()(x)

    # if  not last:
    #     x =  tf.keras.layers.ReLU()(x)
    
    # if use_dropout:
    #     x = layers.Dropout(drop_value)(x)
        
    return x
from tensorflow.keras.layers import SpatialDropout2D
def create_generator():
    inputtensor = Input( shape = (100, ))
    x = Dense(2 * 2 * 512, use_bias=False)(inputtensor)
    # x  = BatchNormalization()(x)
    x =  tf.keras.layers.ReLU()(x)
    x = Reshape( (2,2,512))(x )
    x = gen_block(x, 256 , 1)

    x = gen_block(x, 128 ,2  )#height width  8x8 
    x = gen_block(x, 64 ,3)#height width  16x16
    # x = SpatialDropout2D(0.3)(x)
    x = gen_block(x, 32  , 3, last = True )#height width  32 x 32
    x = Conv2D(3, (1,1 ), strides=1,  padding='same', use_bias=False , activation ='tanh' )(x) # 32 x 32 x 3
    return Model( inputs =  [inputtensor] , outputs = [x ])

# lossgen 
# lossdiscriminator 

# optimgen 

# optimdiscriminator


testgen = create_generator()
testgen.summary()

Model: "model"
_________________________________________________________________
 Layer (type)                Output Shape              Param #   
=================================================================
 input_1 (InputLayer)        [(None, 100)]             0         
                                                                 
 dense (Dense)               (None, 2048)              204800    
                                                                 
 re_lu (ReLU)                (None, 2048)              0         
                                                                 
 reshape (Reshape)           (None, 2, 2, 512)         0         
                                                                 
 up_sampling2d (UpSampling2D  (None, 4, 4, 512)        0         
 )                                                               
                                                                 
 conv2d (Conv2D)             (None, 4, 4, 256)         1179648   
                                                                 
 batch_normalization (BatchN  (None, 4, 4, 256)        1024      
 ormalization)                                                   
                                                                 
 re_lu_1 (ReLU)              (None, 4, 4, 256)         0         
                                                                 
 up_sampling2d_1 (UpSampling  (None, 8, 8, 256)        0         
 2D)                                                             
                                                                 
 conv2d_1 (Conv2D)           (None, 8, 8, 128)         294912    
                                                                 
 batch_normalization_1 (Batc  (None, 8, 8, 128)        512       
 hNormalization)                                                 
                                                                 
 re_lu_2 (ReLU)              (None, 8, 8, 128)         0         
                                                                 
 conv2d_2 (Conv2D)           (None, 8, 8, 128)         147456    
                                                                 
 batch_normalization_2 (Batc  (None, 8, 8, 128)        512       
 hNormalization)                                                 
                                                                 
 re_lu_3 (ReLU)              (None, 8, 8, 128)         0         
                                                                 
 up_sampling2d_2 (UpSampling  (None, 16, 16, 128)      0         
 2D)                                                             
                                                                 
 conv2d_3 (Conv2D)           (None, 16, 16, 64)        73728     
                                                                 
 batch_normalization_3 (Batc  (None, 16, 16, 64)       256       
 hNormalization)                                                 
                                                                 
 re_lu_4 (ReLU)              (None, 16, 16, 64)        0         
                                                                 
 conv2d_4 (Conv2D)           (None, 16, 16, 64)        36864     
                                                                 
 batch_normalization_4 (Batc  (None, 16, 16, 64)       256       
 hNormalization)                                                 
                                                                 
 re_lu_5 (ReLU)              (None, 16, 16, 64)        0         
                                                                 
 conv2d_5 (Conv2D)           (None, 16, 16, 64)        36864     
                                                                 
 batch_normalization_5 (Batc  (None, 16, 16, 64)       256       
 hNormalization)                                                 
                                                                 
 re_lu_6 (ReLU)              (None, 16, 16, 64)        0         
                                                                 
 up_sampling2d_3 (UpSampling  (None, 32, 32, 64)       0         
 2D)                                                             
                                                                 
 conv2d_6 (Conv2D)           (None, 32, 32, 32)        18432     
                                                                 
 batch_normalization_6 (Batc  (None, 32, 32, 32)       128       
 hNormalization)                                                 
                                                                 
 re_lu_7 (ReLU)              (None, 32, 32, 32)        0         
                                                                 
 conv2d_7 (Conv2D)           (None, 32, 32, 32)        9216      
                                                                 
 batch_normalization_7 (Batc  (None, 32, 32, 32)       128       
 hNormalization)                                                 
                                                                 
 re_lu_8 (ReLU)              (None, 32, 32, 32)        0         
                                                                 
 conv2d_8 (Conv2D)           (None, 32, 32, 32)        9216      
                                                                 
 batch_normalization_8 (Batc  (None, 32, 32, 32)       128       
 hNormalization)                                                 
                                                                 
 re_lu_9 (ReLU)              (None, 32, 32, 32)        0         
                                                                 
 conv2d_9 (Conv2D)           (None, 32, 32, 3)         96        
                                                                 
=================================================================
Total params: 2,014,432
Trainable params: 2,012,832
Non-trainable params: 1,600
_________________________________________________________________


createcustomvgg_discriminator().summary()

Model: "model_1"
_________________________________________________________________
 Layer (type)                Output Shape              Param #   
=================================================================
 input_5 (InputLayer)        [(None, 32, 32, 3)]       0         
                                                                 
 conv2d_11 (Conv2D)          (None, 16, 16, 32)        896       
                                                                 
 batch_normalization_10 (Bat  (None, 16, 16, 32)       128       
 chNormalization)                                                
                                                                 
 leaky_re_lu_1 (LeakyReLU)   (None, 16, 16, 32)        0         
                                                                 
 spatial_dropout2d_2 (Spatia  (None, 16, 16, 32)       0         
 lDropout2D)                                                     
                                                                 
 conv2d_12 (Conv2D)          (None, 8, 8, 64)          18496     
                                                                 
 batch_normalization_11 (Bat  (None, 8, 8, 64)         256       
 chNormalization)                                                
                                                                 
 leaky_re_lu_2 (LeakyReLU)   (None, 8, 8, 64)          0         
                                                                 
 spatial_dropout2d_3 (Spatia  (None, 8, 8, 64)         0         
 lDropout2D)                                                     
                                                                 
 conv2d_13 (Conv2D)          (None, 4, 4, 128)         73856     
                                                                 
 batch_normalization_12 (Bat  (None, 4, 4, 128)        512       
 chNormalization)                                                
                                                                 
 leaky_re_lu_3 (LeakyReLU)   (None, 4, 4, 128)         0         
                                                                 
 spatial_dropout2d_4 (Spatia  (None, 4, 4, 128)        0         
 lDropout2D)                                                     
                                                                 
 conv2d_14 (Conv2D)          (None, 2, 2, 256)         295168    
                                                                 
 batch_normalization_13 (Bat  (None, 2, 2, 256)        1024      
 chNormalization)                                                
                                                                 
 leaky_re_lu_4 (LeakyReLU)   (None, 2, 2, 256)         0         
                                                                 
 global_average_pooling2d (G  (None, 256)              0         
 lobalAveragePooling2D)                                          
                                                                 
 dense_1 (Dense)             (None, 1)                 257       
                                                                 
=================================================================
Total params: 390,593
Trainable params: 389,633
Non-trainable params: 960
_________________________________________________________________


from tensorflow.keras.datasets.cifar10 import load_data
wandb.login(key = '')
(x_train, _), (x_val,_) = load_data() 
# Create a GAN object
x_val =(x_val - 127.5) / 127.5
# print(len(x_val))


gan = GAN(discriminator= createcustomvgg_discriminator() , generator=create_generator() , latdim = 100  )

# Specify the optimizers
optimizers = {
    "d_optimizer": tf.keras.optimizers.Adam(learning_rate = 0.0002,beta_1=0.5 ),
    "g_optimizer": tf.keras.optimizers.Adam(learning_rate = 0.0002 , beta_1 =0.5 ),

}

# Specify the losses and metrics
loss = {"d_loss":  tf.keras.losses.BinaryCrossentropy(from_logits=False) , "g_loss": tf.keras.losses.BinaryCrossentropy(from_logits=False)} # trying default binary cross entropy 
# metrics = {"image_loss": metric_fn, "g_loss": None}
# Compile
gan.compile(optimizers=optimizers, loss=loss)
# Start the training and sit back
x_train =(x_train - 127.5) / 127.5
wandb.init(project = 'Vanilla GAN')
import math 
BATCH_SIZE = 128
history = gan.fit(tf.data.Dataset.from_tensor_slices(x_train).batch(BATCH_SIZE), epochs=100 ,   test_step_per_epoch = int(math.ceil((1. * len(x_val)) / BATCH_SIZE)) , validation_data = tf.data.Dataset.from_tensor_slices(x_val).batch(BATCH_SIZE) , callbacks = [
 
    ImageLogger( latent_dim = 100, num_col = 10  , num_row = 5 ) , 
    WandbModelLogger()

    ] )
wandb.finish() 
print( )
# fidhistory = gan.fid_history

wandb: Logging into wandb.ai. (Learn how to deploy a W&B server locally: https://wandb.me/wandb-server)
wandb: You can find your API key in your browser here: https://wandb.ai/authorize
wandb: Paste an API key from your profile and hit enter, or press ctrl+c to quit: wandb: Appending key for api.wandb.ai to your netrc file: /root/.netrc

Downloading data from https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz
170498071/170498071 [==============================] - 6s 0us/step
Downloading data from https://storage.googleapis.com/tensorflow/keras-applications/inception_v3/inception_v3_weights_tf_dim_ordering_tf_kernels_notop.h5
87910968/87910968 [==============================] - 3s 0us/step

wandb: Currently logged in as: yh1234. Use `wandb login --relogin` to force relogin

Epoch 1/100
391/391 [==============================] - ETA: 0s - d_loss: 0.6914 - g_loss: 0.7243

WARNING:tensorflow:Compiled the loaded model, but the compiled metrics have yet to be built. `model.compile_metrics` will be empty until you train or evaluate the model.
WARNING:absl:Found untraced functions such as _jit_compiled_convolution_op, _jit_compiled_convolution_op, _jit_compiled_convolution_op, _jit_compiled_convolution_op while saving (showing 4 of 4). These functions will not be directly callable after loading.
WARNING:tensorflow:Compiled the loaded model, but the compiled metrics have yet to be built. `model.compile_metrics` will be empty until you train or evaluate the model.
WARNING:absl:Found untraced functions such as _jit_compiled_convolution_op, _jit_compiled_convolution_op, _jit_compiled_convolution_op, _jit_compiled_convolution_op, _jit_compiled_convolution_op while saving (showing 5 of 10). These functions will not be directly callable after loading.
wandb: Adding directory to artifact (./Discriminator-Epoch0)... Done. 0.0s
wandb: Adding directory to artifact (./Generator-Epoch0)... Done. 0.0s

391/391 [==============================] - 54s 106ms/step - d_loss: 0.6914 - g_loss: 0.7244 - val_kid: 0.0361
Epoch 2/100
390/391 [============================>.] - ETA: 0s - d_loss: 0.6860 - g_loss: 0.7699

WARNING:tensorflow:Compiled the loaded model, but the compiled metrics have yet to be built. `model.compile_metrics` will be empty until you train or evaluate the model.
WARNING:absl:Found untraced functions such as _jit_compiled_convolution_op, _jit_compiled_convolution_op, _jit_compiled_convolution_op, _jit_compiled_convolution_op while saving (showing 4 of 4). These functions will not be directly callable after loading.
WARNING:tensorflow:Compiled the loaded model, but the compiled metrics have yet to be built. `model.compile_metrics` will be empty until you train or evaluate the model.
WARNING:absl:Found untraced functions such as _jit_compiled_convolution_op, _jit_compiled_convolution_op, _jit_compiled_convolution_op, _jit_compiled_convolution_op, _jit_compiled_convolution_op while saving (showing 5 of 10). These functions will not be directly callable after loading.
wandb: Adding directory to artifact (./Discriminator-Epoch1)... Done. 0.0s
wandb: Adding directory to artifact (./Generator-Epoch1)... Done. 0.0s

391/391 [==============================] - 35s 90ms/step - d_loss: 0.6861 - g_loss: 0.7698 - val_kid: 0.0510
Epoch 3/100
390/391 [============================>.] - ETA: 0s - d_loss: 0.6881 - g_loss: 0.7743

WARNING:tensorflow:Compiled the loaded model, but the compiled metrics have yet to be built. `model.compile_metrics` will be empty until you train or evaluate the model.
WARNING:absl:Found untraced functions such as _jit_compiled_convolution_op, _jit_compiled_convolution_op, _jit_compiled_convolution_op, _jit_compiled_convolution_op while saving (showing 4 of 4). These functions will not be directly callable after loading.
WARNING:tensorflow:Compiled the loaded model, but the compiled metrics have yet to be built. `model.compile_metrics` will be empty until you train or evaluate the model.
WARNING:absl:Found untraced functions such as _jit_compiled_convolution_op, _jit_compiled_convolution_op, _jit_compiled_convolution_op, _jit_compiled_convolution_op, _jit_compiled_convolution_op while saving (showing 5 of 10). These functions will not be directly callable after loading.
wandb: Adding directory to artifact (./Discriminator-Epoch2)... Done. 0.0s
wandb: Adding directory to artifact (./Generator-Epoch2)... Done. 0.0s

391/391 [==============================] - 34s 87ms/step - d_loss: 0.6882 - g_loss: 0.7744 - val_kid: 0.0318
Epoch 4/100
390/391 [============================>.] - ETA: 0s - d_loss: 0.6861 - g_loss: 0.7728


print(history.history)


def plotlearningcurve(history):
    plt.figure(figsize = (16,9))
    plt.subplot(121)
    plt.plot(history.history['d_loss'] , label = 'Discriminator Loss')
    plt.plot(history.history['g_loss'] , label = 'Generator Loss')
    plt.title('Learning Curve - loss')
    plt.ylabel('loss')
    plt.xlabel('Epoch')
    plt.legend( loc='upper left')
    plt.ylim(bottom = 0 )

    plt.subplot(122)
    plt.plot(history.history['val_kid'] , label = 'KID')
    # plt.plot(history.history['val_loss'] , label = 'Validation Accuracy')
    plt.title('Learning Curve - KID')
    plt.ylabel('KID')
    plt.xlabel('Epoch')


    plt.legend( loc='upper left')
    plt.show()


plotlearningcurve(history)


gan.discriminator.save('bcediscriminator.h5')

WARNING:tensorflow:Compiled the loaded model, but the compiled metrics have yet to be built. `model.compile_metrics` will be empty until you train or evaluate the model.


gan.generator.save('bcegenerator.h5')

WARNING:tensorflow:Compiled the loaded model, but the compiled metrics have yet to be built. `model.compile_metrics` will be empty until you train or evaluate the model.


genera


generator = tf.keras.models.load_model('bcegenerator.h5' )

WARNING:tensorflow:No training configuration found in the save file, so the model was *not* compiled. Compile it manually.


import numpy
from numpy import cov
from numpy import trace
from numpy import iscomplexobj
from numpy import asarray
from numpy.random import randint
from scipy.linalg import sqrtm
from keras.applications.inception_v3 import InceptionV3
from keras.applications.inception_v3 import preprocess_input
from keras.datasets.mnist import load_data
# from skimage.transform import resize

def resizeforfid(image):
    image = 255*((image + 1)/2)
    image = tf.image.resize(image, (299, 299), method='nearest')
    image = tf.keras.applications.inception_v3.preprocess_input(image)

    return image
def calculate_fid(act1, act2):
    # calculate mean and covariance statistics
    #https://machinelearningmastery.com/how-to-implement-the-frechet-inception-distance-fid-from-scratch/
    incep= InceptionV3(include_top=False, pooling="avg", input_shape=(299,299,3))
    act1 = incep.predict(act1, verbose=0)
    act2 = incep.predict(act2, verbose=0)
    mu1, sigma1 = act1.mean(axis=0), cov(act1, rowvar=False)
    mu2, sigma2 = act2.mean(axis=0), cov(act2, rowvar=False)
    # calculate sum squared difference between means
    ssdiff = numpy.sum((mu1 - mu2)**2.0)
    # calculate sqrt of product between cov
    covmean = sqrtm(sigma1.dot(sigma2))
    # check and correct imaginary numbers from sqrt
    if iscomplexobj(covmean):          
        covmean = covmean.real
    # calculate score
    fid = ssdiff + trace(sigma1 + sigma2 - 2.0 * covmean)
    return fid

def preparefiddata(realimage, generator):
    realimages = tf.data.Dataset.from_tensor_slices(realimage).map(resizeforfid, num_parallel_calls=tf.data.AUTOTUNE).batch(batch_size=128, num_parallel_calls=tf.data.AUTOTUNE).prefetch(tf.data.AUTOTUNE
    fakeimages = generator(tf.random.normal((len(realimage),100)),training = False)

    fakeimgdataset =tf.data.Dataset.from_tensor_slices(fakeimages).map(resizeforfid, num_parallel_calls=tf.data.AUTOTUNE).batch(batch_size=128, num_parallel_calls=tf.data.AUTOTUNE) .prefetch(tf.data.AUTOTUNE)

    return calculate_fid(realimages, fakeimgdataset)


from tensorflow.keras.datasets.cifar10 import load_data
(x_train, _), (x_val,_) = load_data() 
x_val =(x_val - 127.5) / 127.5
fidscore =preparefiddata(x_val, generator)
print(fidscore)

125.89785529833726

Architecture¶

`Conv2dTranspose` layers are avoided from the start¶

Training result of a GAN with Conv2D Transpose¶

BCE GAN without Conditioning¶

Architecture¶

Conv2dTranspose layers are avoided from the start¶

Training result of a GAN with Conv2D Transpose¶

BCE GAN without Conditioning¶

`Conv2dTranspose` layers are avoided from the start¶